#! /Pub/Users/wangyk/software/pixi/env/omicverse/.pixi/envs/ov/bin/python

import os
import shutil
import glob
import argparse

def organize_and_rename_files(directory: str = ".",
                              suffixes: list = [".barcodes.tsv.gz", ".features.tsv.gz", ".matrix.mtx.gz"]):
    """Organizes files in a directory by grouping files with suffixes '_barcodes.tsv.gz', '_features.tsv.gz', and '_matrix.mtx.gz' into subfolders based on their IDs and renames files in subfolders to remove the ID prefix.

    Args:
        directory (str, optional): The directory to organize files in. Defaults to the current directory.
        suffixes (list, optional): suffixes.
    """

    suffixes = suffixes
    files = glob.glob(os.path.join(directory, "*"))

    file_dict = {}
    for file_path in files:
        filename = os.path.basename(file_path)
        for suffix in suffixes:
            if filename.endswith(suffix):
                base_id = filename[:-len(suffix)]
                if base_id not in file_dict:
                    file_dict[base_id] = {}
                file_dict[base_id][suffix] = file_path
                break

    for base_id, files_with_suffixes in file_dict.items():
        if len(files_with_suffixes) == len(suffixes):
            subfolder_path = os.path.join(directory, base_id)
            os.makedirs(subfolder_path, exist_ok=True)
            for suffix, file_path in files_with_suffixes.items():
                # Remove leading underscore if present for cleaner names
                new_filename = suffix.lstrip('_')
                new_file_path = os.path.join(subfolder_path, new_filename)
                shutil.move(file_path, new_file_path)
                print(
                    f" Organized and renamed {os.path.basename(file_path)} to {new_filename} for ID: {base_id} into folder: {subfolder_path}")
        else:
            print(
                f"Warning: Incomplete set of files for ID: {base_id}. Expected suffixes: {suffixes}, found suffixes: {list(files_with_suffixes.keys())}\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Organize and rename files with suffixes in a directory.')
    parser.add_argument('-d',
        '--directory', type=str, default='.', help='Directory to organize files in.')
    parser.add_argument("-s",
        '--suffixes', nargs='+', default=["barcodes.tsv.gz", "features.tsv.gz", "matrix.mtx.gz"], help='suffixes. default: ["_barcodes.tsv.gz", "_features.tsv.gz", "_matrix.mtx.gz"]')
    args = parser.parse_args()
    
    organize_and_rename_files(args.directory, args.suffixes)
    
    # organize_and_rename_files(
        # "/Pub/Users/wangyk/project/Poroject/F250117001_mr_sc_sepsis/data/")
    # Organizes files in the current directory, you can specify a directory path if needed, e.g., organize_and_rename_files("/path/to/your/directory")
